import numpy as np
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF

def model(X_train, y_train, X_test, y_test):
    fig, axs = plt.subplots(2,3, figsize=(16,7))   
    kernel = RBF(10, (0.01, 1e2))
    TRAINING = 0; TESTING = 1;
    for alpha, column in zip([0.06, 0.8 , 5], range(3)):
        gp = GaussianProcessRegressor(alpha=alpha, kernel = kernel)
        gp.fit(X_train, y_train)
        for set_type, label in zip((TRAINING, TESTING), 
                                   ('Training set', 'Testing set')):
            if set_type == TRAINING:
                # Training accuracy
                training_score = gp.score(X_train, y_train)   
                y_pred = gp.predict(X_train)
                plot(axs, fig, X_train, y_train, y_pred, set_type, column,
                     label+', alpha = '+str(alpha)+', score:
                     {:3.2f}'.format(training_score))
            else:
                # Testing accuracy
                testing_score = gp.score(X_test, y_test)      
                y_pred = gp.predict(X_test)
                plot(axs, fig, X_test, y_test, y_pred, set_type, column,
                     label+', alpha = '+str(alpha)+', score:
                     {:3.2f}'.format(testing_score)) 
    plt.show()
    return

def plot(axs, fig, data2D, target1D, predict1D, set_type, column, label):
    axs[set_type, column].plot(data2D, target1D, 'y.', markersize=12)
    axs[set_type, column].plot(data2D, predict1D, 'k.',markersize=6)
    axs[set_type, column].set_title(label)
    axs[set_type, column].grid()
    return

N = 100  # number of samples                      
X = np.linspace(0,10,N)    
np.random.seed(10)
y = 0.8*X + np.random.randn(N)
X = X[:, np.newaxis]
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=60)
model(X_train, y_train, X_test, y_test)